import numpy as np
import pathlib
import matplotlib.pyplot as plt
import serial
import struct
import threading
import time
from datetime import datetime  # 导入 datetime 模块

# 全局变量，用于控制监听线程的运行
running = True

# 获取父路径
parent_path = pathlib.Path(__file__).parent.absolute()

# 拼接路径
data_path = parent_path / "../data/mnist.npz"
out_path = parent_path / "../../Arm_STM32_FreedomLearn/freedomlearn"

# 解析路径（处理相对路径符号 `..`）
data_path = data_path.resolve()
out_path = out_path.resolve()

# 加载合并的量化权重和比例因子
quantized_params = np.load(f"{parent_path}/quantized_model_parameters.npz")

# 提取量化后的权重和比例因子
w_i_h_quantized = quantized_params['w_i_h_quantized']
w_h_o_quantized = quantized_params['w_h_o_quantized']
b_i_h_quantized = quantized_params['b_i_h_quantized']
b_h_o_quantized = quantized_params['b_h_o_quantized']

w_i_h_scale = quantized_params['w_i_h_scale']
w_h_o_scale = quantized_params['w_h_o_scale']
b_i_h_scale = quantized_params['b_i_h_scale']
b_h_o_scale = quantized_params['b_h_o_scale']

class Dataloader():
    """
        数据读取器
    """  
    def get_data(self):
        with np.load(data_path) as f:
            images, labels = f["x_train"], f["y_train"]
            print("NPZ 文件中的标签（数组名称）:", f.files)
        images = images.astype("uint8") / 255
        images = np.reshape(images, (images.shape[0], images.shape[1] * images.shape[2]))
        labels = np.eye(10)[labels]
        return images, labels

# 反量化函数
def dequantize(data, scale):
    return data * scale / 127.0  # 127 是 int8 的最大值

# 量化推理函数
def quantized_predict(img):
    # 反量化权重
    w_i_h = dequantize(w_i_h_quantized, w_i_h_scale)
    w_h_o = dequantize(w_h_o_quantized, w_h_o_scale)
    b_i_h = dequantize(b_i_h_quantized, b_i_h_scale)
    b_h_o = dequantize(b_h_o_quantized, b_h_o_scale)

    # 前向传播
    h_pre = b_i_h + np.dot(w_i_h, img)
    h = 1 / (1 + np.exp(-h_pre))  # Sigmoid 激活函数
    o_pre = b_h_o + np.dot(w_h_o, h)
    o = 1 / (1 + np.exp(-o_pre))  # Sigmoid 激活函数

    np.set_printoptions(formatter={'float': lambda x: "{:.4f}".format(x)})
    print("输出层的激活值:", o)
    # 返回预测结果
    return np.argmax(o)

# 将图像数据写入 C 文件
def save_image_to_c(file_path, img):
    with open(file_path, "w") as f:
        f.write("float test_image[] = {\n")
        for i, value in enumerate(img.flatten()):
            f.write(f"{value}f, ")
            if (i + 1) % 10 == 0:  # 每行显示 10 个值
                f.write("\n")
        f.write("\n};\n")
    print(f"测试图像已保存到 {file_path}")

def serial_listener(ser):
    """
    持续监听串口并打印接收到的数据。
    """
    while running:
        if ser.in_waiting > 0:  # 检查是否有数据可读
            response = ser.read(ser.in_waiting)  # 读取所有可用的数据
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")  # 获取当前时间戳
            print(f"[{timestamp}] 下位机返回的数据: {response.decode('ascii')}")
        time.sleep(0.1)  # 稍微休眠，避免占用过多CPU

def send_image_over_uart(img, port, baudrate=115200):
    """
    将图像数组转换为字节流并通过串口发送。

    参数:
        img (numpy.ndarray): 图像数组，形状为 (784,)。
        port (str): 串口号（如 'COM3'）。
        baudrate (int): 波特率，默认为 115200。
    """
    global running
    # 打开串口
    try:
        ser = serial.Serial(port, baudrate)
    except serial.SerialException as e:
        print(f"无法打开串口 {port}: {e}")
        return

    # 启动监听线程
    running = True
    listener_thread = threading.Thread(target=serial_listener, args=(ser,))
    listener_thread.start()

    # 将 float 数组转换为字节流
    data = b''
    for value in img:
        data += struct.pack('f', value)  # 将 float 转换为 4 字节

    # 发送数据
    ser.write(data)

    print(f"图像数据已通过串口 {port} 发送")
    # 主线程等待一段时间，确保数据发送完成
    time.sleep(2)

    # 停止监听线程
    running = False
    listener_thread.join()  # 等待监听线程结束
    ser.close()
# 创建数据加载器
dataloader = Dataloader()
images, labels = dataloader.get_data()

# 推理循环
while True:
    index = int(input("输入编号进行预测 (0 - 59999): "))
    img = images[index]
    plt.imshow(img.reshape(28, 28), cmap="Greys")
    img.shape += (1,)

    # 进行量化推理
    predicted_label = quantized_predict(img)

    # 显示预测结果
    plt.title(f"This figure is predicted to be: {predicted_label}")
    plt.show()

    # 通过串口发送图像数据
    send_image_over_uart(img.flatten(), 'COM3')  # 修改为实际的串口号